{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# 实战案例：基于注意力的图像描述模型（ARCTIC）\n",
    "\n",
    "## 模型训练流程\n",
    "\n",
    "![模型训练的一般流程](img/cr-traning_process.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "在现代的深度学习框架基础下，模型训练的一般流程包括读取数据、前馈计算、计算损失、更新参数、选择模型五个步骤。每个步骤需要实现相应的模块。\n",
    "\n",
    "- 在读取数据阶段，我们首先需要下载数据集，然后对整理数据集的格式，以方便接下来构造数据集类，最后在数据集类的基础上构建能够按批次产生训练、验证和测试数据的对象。\n",
    "\n",
    "- 在前馈计算阶段，我们需要实现具体的模型，使得模型能够根据输入产生相应的输出。\n",
    "\n",
    "- 在计算损失阶段，我们需要将模型输出和预期输出进行对比，实现损失函数。\n",
    "\n",
    "- 在更新参数阶段，我们需要给出具体的参数更新方法，即优化方法；由于现代深度学习框架能够自动计算参数梯度，并实现了绝大多数优化方法，我们通常只需要从中进行选择即可。\n",
    "\n",
    "- 在选择模型阶段，我们需要实现具体的评估指标，选出在验证集上表现最优的模型参数。\n",
    "\n",
    "下面，我们将按照这个次序介绍一个图像编码器为CNN网格表示提取器、文本解码器为RNN+注意力的图像描述方法的具体实现。我们的实现大体上是在复现ARCTIC模型，但是在细节上有一些改变，下面的实现过程会对这些改变做具体说明。此外，[链接](https://github.com/sgrvinod/a-PyTorch-Tutorial-to-Image-Captioning)给出了一个更接近原始ARCTIC模型的代码库，非常推荐大家阅读。本节的部分代码也是受到该代码库的启发。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## 读取数据\n",
    "\n",
    "### 下载数据集\n",
    "\n",
    "- 我们使用的数据集为flickr8k([下载地址](https://www.kaggle.com/adityajn105/flickr8k))。\n",
    "- 下载解压后，将其图片放在指定目录(本代码中将该目录设置为../data/flickr8k)下的images文件夹里。\n",
    "- 数据集包括8000张图片，每张图片对应5个句子描述。\n",
    "- 数据集划分采用Karpathy提供的方法([下载地址](http://cs.stanford.edu/people/karpathy/deepimagesent/caption_datasets.zip))，下载解压后，将其中的dataset_flickr8k.json文件拷贝到指定目录下。该划分方法将数据集分成3个子集：6,000张图片和其对应句子描述组成训练集，1,000张图片和描述为验证集，剩余的1,000张图片和描述为测试集。\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### 整理数据集\n",
    "\n",
    "数据集下载完成后，我们需要对其进行处理，以适合之后构造的PyTorch数据集类读取。\n",
    "- 对于文本描述，我们首先构建词典，然后根据词典将文本描述转化为向量。\n",
    "- 对于图像，我们这里仅记录文件路径。\n",
    "- - 如果机器的内存和硬盘空间就比较大，这里也可以将图片读取并处理成三维数组，这样在模型训练和测试的阶段，就不需要再直接读取图片。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "import json\n",
    "import random \n",
    "from collections import defaultdict, Counter\n",
    "from PIL import Image\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "def create_dataset(dataset='flickr8k',\n",
    "                   captions_per_image=5, \n",
    "                   min_word_count=5, \n",
    "                   max_len=30):\n",
    "    \"\"\"\n",
    "    参数：\n",
    "        dataset：数据集名称\n",
    "        captions_per_image：每张图片对应的文本描述数\n",
    "        min_word_count：仅考虑在数据集中（除测试集外）出现5次的词\n",
    "        max_len：文本描述包含的最大单词数，如果文本描述超过该值，则截断\n",
    "    输出：\n",
    "        一个词典文件： vocab.json\n",
    "        三个数据集文件： train_data.json、 val_data.json、 test_data.json\n",
    "    \"\"\"\n",
    "\n",
    "    karpathy_json_path='../data/%s/dataset_flickr8k.json' % dataset\n",
    "    image_folder='../data/%s/images' % dataset\n",
    "    output_folder='../data/%s' % dataset\n",
    "\n",
    "    with open(karpathy_json_path, 'r') as j:\n",
    "        data = json.load(j)\n",
    "    \n",
    "    image_paths = defaultdict(list)\n",
    "    image_captions = defaultdict(list)\n",
    "    vocab = Counter()\n",
    "\n",
    "    for img in data['images']:\n",
    "        split = img['split']\n",
    "        captions = []\n",
    "        for c in img['sentences']:\n",
    "            # 更新词频，测试集在训练过程中是未见数据集，不能统计\n",
    "            if split != 'test':\n",
    "                vocab.update(c['tokens'])\n",
    "            # 不统计超过最大长度限制的词\n",
    "            if len(c['tokens']) <= max_len:\n",
    "                captions.append(c['tokens'])\n",
    "        if len(captions) == 0:\n",
    "            continue\n",
    "\n",
    "        path = os.path.join(image_folder, img['filename'])\n",
    "        \n",
    "        image_paths[split].append(path)\n",
    "        image_captions[split].append(captions)\n",
    "\n",
    "    # 创建词典，增加占位标识符<pad>、未登录词标识符<unk>、句子首尾标识符<start>和<end>\n",
    "    words = [w for w in vocab.keys() if vocab[w] > min_word_count]\n",
    "    vocab = {k: v + 1 for v, k in enumerate(words)}\n",
    "    vocab['<pad>'] = 0\n",
    "    vocab['<unk>'] = len(vocab)\n",
    "    vocab['<start>'] = len(vocab)\n",
    "    vocab['<end>'] = len(vocab)\n",
    "\n",
    "    # 存储词典\n",
    "    with open(os.path.join(output_folder, 'vocab.json'), 'w') as fw:\n",
    "        json.dump(vocab, fw)\n",
    "\n",
    "    # 整理数据集\n",
    "    for split in image_paths:\n",
    "        imgpaths = image_paths[split]\n",
    "        imcaps = image_captions[split]\n",
    "        \n",
    "        enc_captions = []\n",
    "\n",
    "        for i, path in enumerate(imgpaths):\n",
    "            # 合法性检查，检查图像是否可以被解析\n",
    "            img = Image.open(path) \n",
    "            # 如果该图片对应的描述数量不足，则补足\n",
    "            if len(imcaps[i]) < captions_per_image:\n",
    "                captions = imcaps[i] + \\\n",
    "                    [random.choice(imcaps[i]) for _ in range(captions_per_image - len(imcaps[i]))]\n",
    "            # 如果该图片对应的描述数量超了，则随机采样\n",
    "            else:\n",
    "                captions = random.sample(imcaps[i], k=captions_per_image)\n",
    "            assert len(captions) == captions_per_image\n",
    "            \n",
    "            for j, c in enumerate(captions):\n",
    "                # 对文本描述进行编码\n",
    "                enc_c = [vocab['<start>']] + [vocab.get(word, vocab['<unk>']) for word in c] + [vocab['<end>']] \n",
    "                enc_captions.append(enc_c)\n",
    "        # 合法性检查\n",
    "        assert len(imgpaths) * captions_per_image == len(enc_captions)\n",
    "        \n",
    "        # 存储数据\n",
    "        data = {'IMAGES': imgpaths, \n",
    "                'CAPTIONS': enc_captions}\n",
    "        with open(os.path.join(output_folder, split + '_data.json'), 'w') as fw:\n",
    "            json.dump(data, fw)\n",
    "\n",
    "# create_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "在调用该函数生成需要的格式的数据集文件之后，我们可以展示其中一条数据，简单验证下数据的格式是否和我们预想的一致。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 读取词典和验证集\n",
    "with open('../data/flickr8k/vocab.json', 'r') as f:\n",
    "    vocab = json.load(f)\n",
    "vocab_idx2word = {idx:word for word,idx in vocab.items()}\n",
    "with open('../data/flickr8k/val_data.json', 'r') as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "# 展示第12张图片，其对应的文本描述序号是60到64\n",
    "content_img = Image.open(data['IMAGES'][12])\n",
    "plt.imshow(content_img)\n",
    "for i in range(5):\n",
    "    print(' '.join([vocab_idx2word[word_idx] for word_idx in data['CAPTIONS'][12*5+i]]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### 定义数据集类\n",
    "\n",
    "在准备好的数据集的基础上，我们需要进一步定义PyTorch Dataset类，以使用PyTorch DataLoader类按批次产生数据。PyTorch中仅预先定义了图像、文本和语音的单模态任务中常见的数据集类。因此，我们需要定义自己的数据集类。\n",
    "\n",
    "在PyTorch中定义数据集类非常简单，仅需要继承torch.utils.data.Dataset类，并实现\\_\\_getitem\\_\\_和\\_\\_len\\_\\_两个函数即可。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from argparse import Namespace \n",
    "import numpy as np\n",
    "import os\n",
    "import json\n",
    "from PIL import Image\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn.utils.rnn import pack_padded_sequence\n",
    "from torch.utils.data import Dataset\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "class ImageTextDataset(Dataset):\n",
    "    \"\"\"\n",
    "    PyTorch数据类，用于PyTorch DataLoader来按批次产生数据\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dataset_path, vocab_path, split, captions_per_image=5, max_len=30, transform=None):\n",
    "        \"\"\"\n",
    "        参数：\n",
    "            dataset_path：json格式数据文件路径\n",
    "            vocab_path：json格式词典文件路径\n",
    "            split：train、val、test\n",
    "            captions_per_image：每张图片对应的文本描述数\n",
    "            max_len：文本描述包含的最大单词数\n",
    "            transform: 图像预处理方法\n",
    "        \"\"\"\n",
    "        self.split = split\n",
    "        assert self.split in {'train', 'val', 'test'}\n",
    "        self.cpi = captions_per_image\n",
    "        self.max_len = max_len\n",
    "\n",
    "        # 载入数据集\n",
    "        with open(dataset_path, 'r') as f:\n",
    "            self.data = json.load(f)\n",
    "        # 载入词典\n",
    "        with open(vocab_path, 'r') as f:\n",
    "            self.vocab = json.load(f)\n",
    "\n",
    "        # PyTorch图像预处理流程\n",
    "        self.transform = transform\n",
    "\n",
    "        # Total number of datapoints\n",
    "        self.dataset_size = len(self.data['CAPTIONS'])\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        # 第i个文本描述对应第(i // captions_per_image)张图片\n",
    "        img = Image.open(self.data['IMAGES'][i // self.cpi]).convert('RGB')\n",
    "        if self.transform is not None:\n",
    "            img = self.transform(img)\n",
    "\n",
    "        caplen = len(self.data['CAPTIONS'][i])\n",
    "        caption = torch.LongTensor(self.data['CAPTIONS'][i]+ [self.vocab['<pad>']] * (self.max_len + 2 - caplen))\n",
    "        \n",
    "        return img, caption, caplen\n",
    "        \n",
    "\n",
    "    def __len__(self):\n",
    "        return self.dataset_size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### 批量读取数据\n",
    "\n",
    "利用刚才构造的数据集类，借助DataLoader类构建能够按批次产生训练、验证和测试数据的对象。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mktrainval(data_dir, vocab_path, batch_size, workers=4):\n",
    "    train_tx = transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.RandomCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ])\n",
    "    val_tx = transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ])\n",
    "    \n",
    "    train_set = ImageTextDataset(os.path.join(data_dir, 'train_data.json'), \n",
    "                                 vocab_path, 'train',  transform=train_tx)\n",
    "    valid_set = ImageTextDataset(os.path.join(data_dir, 'val_data.json'), \n",
    "                                 vocab_path, 'val', transform=val_tx)\n",
    "    test_set = ImageTextDataset(os.path.join(data_dir, 'test_data.json'), \n",
    "                                 vocab_path, 'test', transform=val_tx)\n",
    "\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        train_set, batch_size=batch_size, shuffle=True, num_workers=workers, pin_memory=True)\n",
    "\n",
    "    valid_loader = torch.utils.data.DataLoader(\n",
    "        valid_set, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=True, drop_last=False)\n",
    "    \n",
    "    test_loader = torch.utils.data.DataLoader(\n",
    "        test_set, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=True, drop_last=False)\n",
    "\n",
    "    return train_loader, valid_loader, test_loader    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## 定义模型\n",
    "\n",
    "ARCTIC模型是一个典型的基于注意力的编解码模型，其编码器为图像网格表示提取器，解码器为循环神经网络。解码器在每生成一个词时，都利用注意力机制考虑当前生成的词和图像中的哪些网格更相关。\n",
    "\n",
    "![ARCTIC的模型结构示意图](img/mt-cnn-attn.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### 图像编码器\n",
    "\n",
    "ARCTIC原始模型使用在ImageNet数据集上预训练过的分类模型VGG19作为图像编码器，VGG19最后一个卷积层作为网格表示提取层。而我们这里使用ResNet-101作为图像编码器，并将其最后一个非全连接层作为网格表示提取层。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import ResNet101_Weights\n",
    "class ImageEncoder(nn.Module):\n",
    "    def __init__(self, finetuned=True):\n",
    "        super(ImageEncoder, self).__init__()\n",
    "        model = torchvision.models.resnet101(weights=ResNet101_Weights.DEFAULT)\n",
    "        # ResNet-101网格表示提取器\n",
    "        self.grid_rep_extractor = nn.Sequential(*(list(model.children())[:-2]))\n",
    "        for param in self.grid_rep_extractor.parameters():\n",
    "            param.requires_grad = finetuned\n",
    "        \n",
    "    def forward(self, images):\n",
    "        out = self.grid_rep_extractor(images) \n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### 文本解码器\n",
    "\n",
    "ARCTIC原始模型使用结合注意力的LSTM作为文本解码器，我们这里使用结合注意力的GRU作为文本解码器，注意力评分函数采用的是加性注意力。下面给出加性注意力和解码器的具体实现。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "加性注意力评分函数的具体形式为 $W_2^T{\\rm tanh}(W_1 [\\mathbf{q}_i; \\mathbf{k}_j])$ 。\n",
    "\n",
    "- 首先将权重 $W_1$ 依照查询q和键k的维度，相应地拆成两组权重，分别将单个查询和一组键映射到到注意力函数隐藏层表示空间；\n",
    "- 然后将二者相加得到一组维度为attn_dim的表示，并在经过非线性变换后，使用形状为(attn_dim, 1) 的权重 $W_2$ 将其映射为一组数值；\n",
    "- 再通过softmax函数获取单个查询和所有键的关联程度，即归一化的相关性分数；\n",
    "- 最后以相关性得分为权重，对值进行加权求和，计算输出特征。这里的值和键是同一组向量表示。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AdditiveAttention(nn.Module):\n",
    "    def  __init__(self, query_dim, key_dim, attn_dim):\n",
    "        \"\"\"\n",
    "        参数：\n",
    "            query_dim: 查询Q的维度\n",
    "            key_dim: 键K的维度\n",
    "            attn_dim: 注意力函数隐藏层表示的维度\n",
    "        \"\"\"\n",
    "        super(AdditiveAttention, self).__init__()\n",
    "        self.attn_w_1_q = nn.Linear(query_dim, attn_dim)\n",
    "        self.attn_w_1_k = nn.Linear(key_dim, attn_dim)\n",
    "        self.attn_w_2 = nn.Linear(attn_dim, 1)\n",
    "        self.tanh = nn.Tanh()\n",
    "        self.softmax = nn.Softmax(dim=1)\n",
    "\n",
    "    def forward(self, query, key_value):\n",
    "        \"\"\"\n",
    "        Q K V：Q和K算出相关性得分，作为V的权重，K=V\n",
    "        参数：\n",
    "            query: 查询 (batch_size, q_dim)\n",
    "            key_value: 键和值，(batch_size, n_kv, kv_dim)\n",
    "        \"\"\"\n",
    "        # （2）计算query和key的相关性，实现注意力评分函数\n",
    "        # -> (batch_size, 1, attn_dim)\n",
    "        queries = self.attn_w_1_q(query).unsqueeze(1)\n",
    "        # -> (batch_size, n_kv, attn_dim)\n",
    "        keys = self.attn_w_1_k(key_value)\n",
    "        # -> (batch_size, n_kv)\n",
    "        attn = self.attn_w_2(self.tanh(queries+keys)).squeeze(2) \n",
    "        # （3）归一化相关性分数\n",
    "        # -> (batch_size, n_kv)\n",
    "        attn = self.softmax(attn) \n",
    "        # （4）计算输出\n",
    "        # (batch_size x 1 x n_kv)(batch_size x n_kv x kv_dim)\n",
    "        # -> (batch_size, 1, kv_dim)\n",
    "        output = torch.bmm(attn.unsqueeze(1), key_value).squeeze(1)\n",
    "        return output, attn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "解码器前馈过程的实现流程如下：\n",
    "\n",
    "（1）将图文数据按照文本的实际长度从长到短排序，这是为了模拟pack_padded_sequence函数的思想，方便后面使用动态的批大小，以避免<pad>参与运算带来的非必要的计算消耗。\n",
    "\n",
    "![pack_padded_sequence函数的作用的示例图](img/cr-pack_padded_sequence-example.png)\n",
    "    \n",
    "\n",
    "（2）在第一时刻解码前，使用图像表示来初始化GRU的隐状态。\n",
    "\n",
    "（3）解码的每一时刻的具体操作可以分解为如下4个子操作：\n",
    "    \n",
    "- （3.1）获取实际的批大小；\n",
    "\n",
    "- （3.2）利用GRU前一时刻最后一个隐藏层的状态作为查询，图像表示作为键和值，获取上下文向量；\n",
    "\n",
    "- （3.3）将上下文向量和当前时刻输入的词表示拼接起来，作为GRU该时刻的输入，获得输出；\n",
    "\n",
    "- （3.4）使用全连接层和softmax激活函数将GRU的输出映射为词表上的概率分布。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionDecoder(nn.Module):\n",
    "    def __init__(self, image_code_dim, vocab_size, word_dim, attention_dim, hidden_size, num_layers, dropout=0.5):\n",
    "        super(AttentionDecoder, self).__init__()\n",
    "        self.embed = nn.Embedding(vocab_size, word_dim)\n",
    "        self.attention = AdditiveAttention(hidden_size, image_code_dim, attention_dim)\n",
    "        self.init_state = nn.Linear(image_code_dim, num_layers*hidden_size)\n",
    "        self.rnn = nn.GRU(word_dim + image_code_dim, hidden_size, num_layers)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        self.fc = nn.Linear(hidden_size, vocab_size)\n",
    "        # RNN默认已初始化\n",
    "        self.init_weights()\n",
    "        \n",
    "    def init_weights(self):\n",
    "        self.embed.weight.data.uniform_(-0.1, 0.1)\n",
    "        self.fc.bias.data.fill_(0)\n",
    "        self.fc.weight.data.uniform_(-0.1, 0.1)\n",
    "    \n",
    "    def init_hidden_state(self, image_code, captions, cap_lens):\n",
    "        \"\"\"\n",
    "        参数：\n",
    "            image_code：图像编码器输出的图像表示 \n",
    "                        (batch_size, image_code_dim, grid_height, grid_width)\n",
    "        \"\"\"\n",
    "        # 将图像网格表示转换为序列表示形式 \n",
    "        batch_size, image_code_dim = image_code.size(0), image_code.size(1)\n",
    "        # -> (batch_size, grid_height, grid_width, image_code_dim) \n",
    "        image_code = image_code.permute(0, 2, 3, 1)  \n",
    "        # -> (batch_size, grid_height * grid_width, image_code_dim)\n",
    "        image_code = image_code.view(batch_size, -1, image_code_dim)\n",
    "        # （1）按照caption的长短排序\n",
    "        sorted_cap_lens, sorted_cap_indices = torch.sort(cap_lens, 0, True)\n",
    "        captions = captions[sorted_cap_indices]\n",
    "        image_code = image_code[sorted_cap_indices]\n",
    "         #（2）初始化隐状态\n",
    "        hidden_state = self.init_state(image_code.mean(axis=1))\n",
    "        hidden_state = hidden_state.view(\n",
    "                            batch_size, \n",
    "                            self.rnn.num_layers, \n",
    "                            self.rnn.hidden_size).permute(1, 0, 2)\n",
    "        return image_code, captions, sorted_cap_lens, sorted_cap_indices, hidden_state\n",
    "\n",
    "    def forward_step(self, image_code, curr_cap_embed, hidden_state):\n",
    "        #（3.2）利用注意力机制获得上下文向量\n",
    "        # query：hidden_state[-1]，即最后一个隐藏层输出 (batch_size, hidden_size)\n",
    "        # context: (batch_size, hidden_size)\n",
    "        context, alpha = self.attention(hidden_state[-1], image_code)\n",
    "        #（3.3）以上下文向量和当前时刻词表示为输入，获得GRU输出\n",
    "        x = torch.cat((context, curr_cap_embed), dim=-1).unsqueeze(0)\n",
    "        # x: (1, real_batch_size, hidden_size+word_dim)\n",
    "        # out: (1, real_batch_size, hidden_size)\n",
    "        out, hidden_state = self.rnn(x, hidden_state)\n",
    "        #（3.4）获取该时刻的预测结果\n",
    "        # (real_batch_size, vocab_size)\n",
    "        preds = self.fc(self.dropout(out.squeeze(0)))\n",
    "        return preds, alpha, hidden_state\n",
    "        \n",
    "    def forward(self, image_code, captions, cap_lens):\n",
    "        \"\"\"\n",
    "        参数：\n",
    "            hidden_state: (num_layers, batch_size, hidden_size)\n",
    "            image_code:  (batch_size, feature_channel, feature_size)\n",
    "            captions: (batch_size, )\n",
    "        \"\"\"\n",
    "        # （1）将图文数据按照文本的实际长度从长到短排序\n",
    "        # （2）获得GRU的初始隐状态\n",
    "        image_code, captions, sorted_cap_lens, sorted_cap_indices, hidden_state \\\n",
    "            = self.init_hidden_state(image_code, captions, cap_lens)\n",
    "        batch_size = image_code.size(0)\n",
    "        # 输入序列长度减1，因为最后一个时刻不需要预测下一个词\n",
    "        lengths = sorted_cap_lens.cpu().numpy() - 1\n",
    "        # 初始化变量：模型的预测结果和注意力分数\n",
    "        predictions = torch.zeros(batch_size, lengths[0], self.fc.out_features).to(captions.device)\n",
    "        alphas = torch.zeros(batch_size, lengths[0], image_code.shape[1]).to(captions.device)\n",
    "        # 获取文本嵌入表示 cap_embeds: (batch_size, num_steps, word_dim)\n",
    "        cap_embeds = self.embed(captions)\n",
    "        # Teacher-Forcing模式\n",
    "        for step in range(lengths[0]):\n",
    "            #（3）解码\n",
    "            #（3.1）模拟pack_padded_sequence函数的原理，获取该时刻的非<pad>输入\n",
    "            real_batch_size = np.where(lengths>step)[0].shape[0]\n",
    "            preds, alpha, hidden_state = self.forward_step(\n",
    "                            image_code[:real_batch_size], \n",
    "                            cap_embeds[:real_batch_size, step, :],\n",
    "                            hidden_state[:, :real_batch_size, :].contiguous())            \n",
    "            # 记录结果\n",
    "            predictions[:real_batch_size, step, :] = preds\n",
    "            alphas[:real_batch_size, step, :] = alpha\n",
    "        return predictions, alphas, captions, lengths, sorted_cap_indices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### ARCTIC模型\n",
    "\n",
    "在定义编码器和解码器完成之后，我们就很容易构建图像描述模型ARCTIC了。仅需要在初始化函数时声明编码器和解码器，然后在前馈函数实现里，将编码器的输出和文本描述作为解码器的输入即可。\n",
    "\n",
    "这里我们额外定义了束搜索采样函数，用于生成句子，以计算BLEU值。下面的代码详细标注了其具体实现。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ARCTIC(nn.Module):\n",
    "    def __init__(self, image_code_dim, vocab, word_dim, attention_dim, hidden_size, num_layers):\n",
    "        super(ARCTIC, self).__init__()\n",
    "        self.vocab = vocab\n",
    "        self.encoder = ImageEncoder()\n",
    "        self.decoder = AttentionDecoder(image_code_dim, len(vocab), word_dim, attention_dim, hidden_size, num_layers)\n",
    "\n",
    "    def forward(self, images, captions, cap_lens):\n",
    "        image_code = self.encoder(images)\n",
    "        return self.decoder(image_code, captions, cap_lens)\n",
    "    \n",
    "    def generate_by_beamsearch(self, images, beam_k, max_len):\n",
    "        vocab_size = len(self.vocab)\n",
    "        image_codes = self.encoder(images)\n",
    "        texts = []\n",
    "        device = images.device\n",
    "        # 对每个图像样本执行束搜索\n",
    "        for image_code in image_codes:\n",
    "            # 将图像表示复制k份\n",
    "            image_code = image_code.unsqueeze(0).repeat(beam_k,1,1,1)\n",
    "            # 生成k个候选句子，初始时，仅包含开始符号<start>\n",
    "            cur_sents = torch.full((beam_k, 1), self.vocab['<start>'], dtype=torch.long).to(device)\n",
    "            cur_sent_embed = self.decoder.embed(cur_sents)[:,0,:]\n",
    "            sent_lens = torch.LongTensor([1]*beam_k).to(device)\n",
    "            # 获得GRU的初始隐状态\n",
    "            image_code, cur_sent_embed, _, _, hidden_state = \\\n",
    "                self.decoder.init_hidden_state(image_code, cur_sent_embed, sent_lens)\n",
    "            # 存储已生成完整的句子（以句子结束符<end>结尾的句子）\n",
    "            end_sents = []\n",
    "            # 存储已生成完整的句子的概率\n",
    "            end_probs = []\n",
    "            # 存储未完整生成的句子的概率\n",
    "            probs = torch.zeros(beam_k, 1).to(device)\n",
    "            k = beam_k\n",
    "            while True:\n",
    "                preds, _, hidden_state = self.decoder.forward_step(image_code[:k], cur_sent_embed, hidden_state.contiguous())\n",
    "                # -> (k, vocab_size)\n",
    "                preds = nn.functional.log_softmax(preds, dim=1)\n",
    "                # 对每个候选句子采样概率值最大的前k个单词生成k个新的候选句子，并计算概率\n",
    "                # -> (k, vocab_size)\n",
    "                probs = probs.repeat(1,preds.size(1)) + preds\n",
    "                if cur_sents.size(1) == 1:\n",
    "                    # 第一步时，所有句子都只包含开始标识符，因此，仅利用其中一个句子计算topk\n",
    "                    values, indices = probs[0].topk(k, 0, True, True)\n",
    "                else:\n",
    "                    # probs: (k, vocab_size) 是二维张量\n",
    "                    # topk函数直接应用于二维张量会按照指定维度取最大值，这里需要在全局取最大值\n",
    "                    # 因此，将probs转换为一维张量，再使用topk函数获取最大的k个值\n",
    "                    values, indices = probs.view(-1).topk(k, 0, True, True)\n",
    "                # 计算最大的k个值对应的句子索引和词索引\n",
    "                sent_indices = torch.div(indices, vocab_size, rounding_mode='trunc') \n",
    "                word_indices = indices % vocab_size \n",
    "                # 将词拼接在前一轮的句子后，获得此轮的句子\n",
    "                cur_sents = torch.cat([cur_sents[sent_indices], word_indices.unsqueeze(1)], dim=1)\n",
    "                # 查找此轮生成句子结束符<end>的句子\n",
    "                end_indices = [idx for idx, word in enumerate(word_indices) if word == self.vocab['<end>']]\n",
    "                if len(end_indices) > 0:\n",
    "                    end_probs.extend(values[end_indices])\n",
    "                    end_sents.extend(cur_sents[end_indices].tolist())\n",
    "                    # 如果所有的句子都包含结束符，则停止生成\n",
    "                    k -= len(end_indices)\n",
    "                    if k == 0:\n",
    "                        break\n",
    "                # 查找还需要继续生成词的句子\n",
    "                cur_indices = [idx for idx, word in enumerate(word_indices) \n",
    "                               if word != self.vocab['<end>']]\n",
    "                if len(cur_indices) > 0:\n",
    "                    cur_sent_indices = sent_indices[cur_indices]\n",
    "                    cur_word_indices = word_indices[cur_indices]\n",
    "                    # 仅保留还需要继续生成的句子、句子概率、隐状态、词嵌入\n",
    "                    cur_sents = cur_sents[cur_indices]\n",
    "                    probs = values[cur_indices].view(-1,1)\n",
    "                    hidden_state = hidden_state[:,cur_sent_indices,:]\n",
    "                    cur_sent_embed = self.decoder.embed(\n",
    "                        cur_word_indices.view(-1,1))[:,0,:]\n",
    "                # 句子太长，停止生成\n",
    "                if cur_sents.size(1) >= max_len:\n",
    "                    break\n",
    "            if len(end_sents) == 0:\n",
    "                # 如果没有包含结束符的句子，则选取第一个句子作为生成句子\n",
    "                gen_sent = cur_sents[0].tolist()\n",
    "            else: \n",
    "                # 否则选取包含结束符的句子中概率最大的句子\n",
    "                gen_sent = end_sents[end_probs.index(max(end_probs))]\n",
    "            texts.append(gen_sent)\n",
    "        return texts\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## 定义损失函数\n",
    "\n",
    "这里采用了最常用的交叉熵损失作为损失函数。由于同一个训练批次里的文本描述的长度不一致，因此，有大量的不需要计算损失的<pad>目标。为了避免计算资源的浪费，这里先将数据按照文本长度排序，再利用pack_padded_sequence函数将预测目标为\\<pad\\>的数据去除，最后再利用交叉熵损失计算实际的损失。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PackedCrossEntropyLoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(PackedCrossEntropyLoss, self).__init__()\n",
    "        self.loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, predictions, targets, lengths):\n",
    "        \"\"\"\n",
    "        参数：\n",
    "            predictions：按文本长度排序过的预测结果\n",
    "            targets：按文本长度排序过的文本描述\n",
    "            lengths：文本长度\n",
    "        \"\"\"\n",
    "        predictions = pack_padded_sequence(predictions, lengths, batch_first=True)[0]\n",
    "        targets = pack_padded_sequence(targets, lengths, batch_first=True)[0]\n",
    "        return self.loss_fn(predictions, targets)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## 选择优化方法\n",
    "\n",
    "这里选用Adam优化算法来更新模型参数，由于数据集较小，训练轮次少，因此，学习速率在训练过程中并不调整。但是对编码器和解码器采用了不同的学习速率。具体来说，预训练的图像编码器的学习速率小于需要从头开始训练的文本解码器的学习速率。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_optimizer(model, config):\n",
    "    return torch.optim.Adam([{\"params\": filter(lambda p: p.requires_grad, model.encoder.parameters()), \n",
    "                              \"lr\": config.encoder_learning_rate},\n",
    "                             {\"params\": filter(lambda p: p.requires_grad, model.decoder.parameters()), \n",
    "                              \"lr\": config.decoder_learning_rate}])\n",
    "    \n",
    "def adjust_learning_rate(optimizer, epoch, config):\n",
    "    \"\"\"\n",
    "        每隔lr_update个轮次，学习速率减小至当前十分之一，\n",
    "        实际上，我们并未使用该函数，这里是为了展示在训练过程中调整学习速率的方法。\n",
    "    \"\"\"\n",
    "    optimizer.param_groups[0]['lr'] = config.encoder_learning_rate * (0.1 ** (epoch // config.lr_update))\n",
    "    optimizer.param_groups[1]['lr'] = config.decoder_learning_rate * (0.1 ** (epoch // config.lr_update))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## 评估指标\n",
    "\n",
    "这里借助nltk库实现了图像描述中最常用的评估指标BLEU值，需要注意的是，再调用计算BLEU值之前，要先将文本中人工添加的文本开始符、结束符和占位符去掉。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nltk.translate.bleu_score import corpus_bleu\n",
    "\n",
    "def filter_useless_words(sent, filterd_words):\n",
    "    # 去除句子中不参与BLEU值计算的符号\n",
    "    return [w for w in sent if w not in filterd_words]\n",
    "\n",
    "def evaluate(data_loader, model, config):\n",
    "    model.eval()\n",
    "    # 存储候选文本\n",
    "    cands = []\n",
    "    # 存储参考文本\n",
    "    refs = []\n",
    "    # 需要过滤的词\n",
    "    filterd_words = set({model.vocab['<start>'], model.vocab['<end>'], model.vocab['<pad>']})\n",
    "    cpi = config.captions_per_image\n",
    "    device = next(model.parameters()).device\n",
    "    for i, (imgs, caps, caplens) in enumerate(data_loader):\n",
    "        with torch.no_grad():\n",
    "            # 通过束搜索，生成候选文本\n",
    "            texts = model.generate_by_beamsearch(imgs.to(device), config.beam_k, config.max_len+2)\n",
    "            # 候选文本\n",
    "            cands.extend([filter_useless_words(text, filterd_words) for text in texts])\n",
    "            # 参考文本\n",
    "            refs.extend([filter_useless_words(cap, filterd_words) for cap in caps.tolist()])\n",
    "    # 实际上，每个候选文本对应cpi条参考文本\n",
    "    multiple_refs = []\n",
    "    for idx in range(len(refs)):\n",
    "        multiple_refs.append(refs[(idx//cpi)*cpi : (idx//cpi)*cpi+cpi])\n",
    "    # 计算BLEU-4值，corpus_bleu函数默认weights权重为(0.25,0.25,0.25,0.25)\n",
    "    # 即计算1-gram到4-gram的BLEU几何平均值\n",
    "    bleu4 = corpus_bleu(multiple_refs, cands, weights=(0.25,0.25,0.25,0.25))\n",
    "    model.train()\n",
    "    return bleu4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## 训练模型\n",
    "\n",
    "训练模型过程仍然是分为读取数据、前馈计算、计算损失、更新参数、选择模型五个步骤。\n",
    "\n",
    "模型训练的具体方案为一共训练30轮，编码器和解码器的学习速率分别为0.0001和0.0005。\n",
    "\n",
    "<!-- 模型训练的具体方案为一共训练30轮，初始编码器和解码器的学习速率分别为0.0001和0.0005，每10轮将学习速率变为原数值的1/10。 -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置模型超参数和辅助变量\n",
    "config = Namespace(\n",
    "    max_len = 30,\n",
    "    captions_per_image = 5,\n",
    "    batch_size = 32,\n",
    "    image_code_dim = 2048,\n",
    "    word_dim = 512,\n",
    "    hidden_size = 512,\n",
    "    attention_dim = 512,\n",
    "    num_layers = 1,\n",
    "    encoder_learning_rate = 0.0001,\n",
    "    decoder_learning_rate = 0.0005,\n",
    "    num_epochs = 10,\n",
    "    grad_clip = 5.0,\n",
    "    alpha_weight = 1.0,\n",
    "    evaluate_step = 900, # 每隔多少步在验证集上测试一次\n",
    "    checkpoint = None, # 如果不为None，则利用该变量路径的模型继续训练\n",
    "    best_checkpoint = '../model/ARCTIC/best_flickr8k.ckpt', # 验证集上表现最优的模型的路径\n",
    "    last_checkpoint = '../model/ARCTIC/last_flickr8k.ckpt', # 训练完成时的模型的路径\n",
    "    beam_k = 5\n",
    ")\n",
    "\n",
    "# 设置GPU信息\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '2'\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\") \n",
    "\n",
    "# 数据\n",
    "data_dir = '../data/flickr8k/'\n",
    "vocab_path = '../data/flickr8k/vocab.json'\n",
    "train_loader, valid_loader, test_loader = mktrainval(data_dir, vocab_path, config.batch_size)\n",
    "\n",
    "# 模型\n",
    "with open(vocab_path, 'r') as f:\n",
    "    vocab = json.load(f)\n",
    "\n",
    "# 随机初始化 或 载入已训练的模型\n",
    "start_epoch = 0\n",
    "checkpoint = config.checkpoint\n",
    "if checkpoint is None:\n",
    "    model = ARCTIC(config.image_code_dim, vocab, config.word_dim, config.attention_dim, config.hidden_size, config.num_layers)\n",
    "else:\n",
    "    checkpoint = torch.load(checkpoint)\n",
    "    start_epoch = checkpoint['epoch'] + 1\n",
    "    model = checkpoint['model']\n",
    "\n",
    "# 优化器\n",
    "optimizer = get_optimizer(model, config)\n",
    "\n",
    "# 将模型拷贝至GPU，并开启训练模式\n",
    "model.to(device)\n",
    "model.train()\n",
    "\n",
    "# 损失函数\n",
    "loss_fn = PackedCrossEntropyLoss().to(device)\n",
    "\n",
    "best_res = 0\n",
    "print(\"开始训练\")\n",
    "fw = open('log.txt', 'w')\n",
    "\n",
    "for epoch in range(start_epoch, config.num_epochs):\n",
    "    for i, (imgs, caps, caplens) in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        # 1. 读取数据至GPU\n",
    "        imgs = imgs.to(device)\n",
    "        caps = caps.to(device)\n",
    "        caplens = caplens.to(device)\n",
    "\n",
    "        # 2. 前馈计算\n",
    "        predictions, alphas, sorted_captions, lengths, sorted_cap_indices = model(imgs, caps, caplens)\n",
    "        # 3. 计算损失\n",
    "        # captions从第2个词开始为targets\n",
    "        loss = loss_fn(predictions, sorted_captions[:, 1:], lengths)\n",
    "        # 重随机注意力正则项，使得模型尽可能全面的利用到每个网格\n",
    "        # 要求所有时刻在同一个网格上的注意力分数的平方和接近1\n",
    "        loss += config.alpha_weight * ((1. - alphas.sum(axis=1)) ** 2).mean()\n",
    "\n",
    "        loss.backward()\n",
    "        # 梯度截断\n",
    "        if config.grad_clip > 0:\n",
    "            nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)\n",
    "        \n",
    "        # 4. 更新参数\n",
    "        optimizer.step()\n",
    "        \n",
    "        if (i+1) % 100 == 0:\n",
    "            print('epoch %d, step %d: loss=%.2f' % (epoch, i+1, loss.cpu()))\n",
    "            fw.write('epoch %d, step %d: loss=%.2f \\n' % (epoch, i+1, loss.cpu()))\n",
    "            fw.flush()\n",
    "\n",
    "        state = {\n",
    "                'epoch': epoch,\n",
    "                'step': i,\n",
    "                'model': model,\n",
    "                'optimizer': optimizer\n",
    "                }\n",
    "        if (i+1) % config.evaluate_step == 0:\n",
    "            bleu_score = evaluate(valid_loader, model, config)\n",
    "            # 5. 选择模型\n",
    "            if best_res < bleu_score:\n",
    "                best_res = bleu_score\n",
    "                torch.save(state, config.best_checkpoint)\n",
    "            torch.save(state, config.last_checkpoint)\n",
    "            fw.write('Validation@epoch, %d, step, %d, BLEU-4=%.2f\\n' % \n",
    "                  (epoch, i+1, bleu_score))\n",
    "            fw.flush()\n",
    "            print('Validation@epoch, %d, step, %d, BLEU-4=%.2f' % \n",
    "                  (epoch, i+1, bleu_score))\n",
    "checkpoint = torch.load(config.best_checkpoint)\n",
    "model = checkpoint['model']\n",
    "bleu_score = evaluate(test_loader, model, config)\n",
    "print(\"Evaluate on the test set with the model that has the best performance on the validation set\")\n",
    "print('Epoch: %d, BLEU-4=%.2f' % \n",
    "      (checkpoint['epoch'], bleu_score))\n",
    "fw.write('Epoch: %d, BLEU-4=%.2f' % \n",
    "      (checkpoint['epoch'], bleu_score))\n",
    "fw.close()      \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这段代码完成训练，最后一行会输出在验证集上表现最好的模型在测试集上的结果，具体如下：\n",
    "\n",
    "```\n",
    "Epoch: 4, BLEU-4=0.23\n",
    "```\n"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "rise": {
   "enable_chalkboard": true,
   "scroll": true
  },
  "vscode": {
   "interpreter": {
    "hash": "9c90bab175f6ebea40ddbaa491ed6726e7c53d0208272849c54b7fb296566795"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
